from collections import Counter

from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.model_selection import ShuffleSplit


def test01():
    # 加载数据
    x, y = load_iris(return_X_y=True)
    # 留出法（随机分割）
    x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2)
    print('随机类别分割:', Counter(y_train), Counter(y_test))
    # 留出法 分层分割
    x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2, stratify=y)
    print('分层分割别分割:', Counter(y_train), Counter(y_test))

def test02():
    x, y = load_iris(return_X_y=True)
    print(Counter(y))
    # 留出法 随机分割
    splitter = ShuffleSplit(n_splits=5, test_size=0.2, random_state=42)
    for train, test in splitter.split(x, y):
        print('随机分割:', Counter(y[test]))

    # 留出法 分层分割
    from sklearn.model_selection import StratifiedShuffleSplit
    splitter = StratifiedShuffleSplit(n_splits=5, test_size=0.2, random_state=42)
    for train, test in splitter.split(x, y):
        print('分层分割:', Counter(y[test]))

if __name__ == '__main__':
    test02()